# -*- coding: utf-8 -*-
import numpy as np
from node import Node
from ElementBase import elementBase
from tools import *
from typing import List
from Loads import *
# 框架单元类

class F3D2_EB(elementBase):
    """
    空间框架单元,EB后缀代表采用欧拉伯努利梁假设,节点自由度:UX,UY,UZ,URX,URY,URZ;
    
    可承受载荷: 轴力,弯矩,扭矩,不支持单元载荷.
    
    单元定义需要:A,E,nu,rho,节点坐标,Iyy,Izz,J;(Iyy,Izz为截面惯性矩,J为截面极惯性矩);
    
    单元公式见项目文档;
    """
    def __init__(self, label:int,
                    node1:Node,
                    node2:Node,
                    E=0.0,nu=0.3,A=0.0,Iyy=0.0,Izz=0.0,rho=0.0,J=0.0):
        super().__init__(label, 6)
        self.node1=node1
        self.node2=node2
        self.E=E
        self.nu=nu  # 泊松比
        self.A=A
        self.Izz=Izz
        self.Iyy=Iyy
        self.J=J  # 截面极惯性矩(iint_A{r^2 dA})
        self.rho=rho
        
        # 其他设置
        # 默认局部坐标轴,使用SAP2000的局部坐标系定义:
        # 1. 若局部x轴和全局Z轴平行,则局部y轴为全局+X轴,局部z轴=cross(局部x,局部y);
        # 2. 若局部x轴和全局Z轴不平行,则局部z轴=cross(局部x,全局Z),局部y轴=cross(局部z,局部x);
        self.local_orient=np.zeros((3,3),dtype=float)
        # 局部x轴方向
        self.local_orient[0,:]=(self.node2.coordinate-self.node1.coordinate)/self.L
        if np.allclose(self.local_orient[0,:], [0, 0, 1]) or np.allclose(self.local_orient[0,:], [0, 0, -1]):
            # 局部x轴和全局Z轴平行
            self.local_orient[1,:]=np.array([1.0, 0, 0])
            self.local_orient[2,:]=np.cross(self.local_orient[0,:], self.local_orient[1,:])
        else:
            # 局部x轴和全局Z轴不平行
            self.local_orient[2,:]=np.cross(self.local_orient[0,:], np.array([0, 0, 1]))
            self.local_orient[1,:]=np.cross(self.local_orient[2,:], self.local_orient[0,:])
        # 单位化局部坐标系
        self.local_orient /= np.linalg.norm(self.local_orient, axis=1)[:, np.newaxis]
    @property
    def get_connectivity(self):
        """
        获取单元节点编号
        
        returns:
            List[int]: 节点编号列表
        """
        return [self.node1.label, self.node2.label]
    @property
    def L(self):
        """
        计算单元长度
        
        returns:
            float: 单元长度
        """
        return np.linalg.norm(self.node2.coordinate - self.node1.coordinate)
    @property
    def t(self):
        """
        计算变换矩阵[lamda]
        
        returns:
            np.ndarray: 变换矩阵(12,12)
        """
        # 局部坐标系矩阵,应当是单位向量
        t=np.zeros((12,12),dtype=float)
        t[0:3,0:3]=self.local_orient
        t[3:6,3:6]=self.local_orient
        t[6:9,6:9]=self.local_orient
        t[9:12,9:12]=self.local_orient
        return t
    @property
    def Ke(self):
        """
        计算全局坐标系下的单元刚度矩阵
        
        returns:
            np.ndarray: 单元刚度矩阵(12,12)
        """
        ke=np.zeros((12,12),dtype=float)
        # 装配 Axial Displacements 的刚度矩阵部分
        ke[0,0]+=(self.E*self.A)/self.L
        ke[6,6]+=ke[0,0]
        ke[0,6]+=-1*(self.E*self.A)/self.L
        ke[6,0]+=ke[0,6]
        # 加入 Torsional Displacements 的刚度矩阵部分
        G=self.E/(2*(1+self.nu))  # 剪切模量
        ke[9,9]=self.J*G/self.L
        ke[3,3]+=ke[9,9]
        ke[9,3]=-1.0*self.J*G/self.L
        ke[3,9]+=ke[9,3]
        # 加入 Bending Displacements in the Plane xy 的刚度矩阵部分
        kxy=np.array(([[12, 6*self.L,-12, 6*self.L],
                        [6*self.L, 4*self.L**2,-6*self.L, 2*self.L**2],
                        [-12,-6*self.L,12,-6*self.L],
                        [6*self.L, 2*self.L**2,-6*self.L, 4*self.L**2]]))*(self.E*self.Izz/(self.L**3))
        ke[np.ix_([1,5,7,11],[1,5,7,11])]+=kxy
        # 加入 Bending Displacements in the Plane xz 的刚度矩阵部分
        kxz=np.array(([[12, 6*self.L,-12, 6*self.L],
                        [6*self.L, 4*self.L**2,-6*self.L, 2*self.L**2],
                        [-12,-6*self.L,12,-6*self.L],
                        [6*self.L, 2*self.L**2,-6*self.L, 4*self.L**2]]))*(self.E*self.Iyy/(self.L**3))
        ke[np.ix_([2,4,8,10],[2,4,8,10])]+=kxz
        return self.t.T@ke@self.t
    @property
    def Me(self):
        """
        计算全局坐标系下的单元质量矩阵(集中质量矩阵)
        
        returns:
            np.ndarray: 单元质量矩阵(12,12)
        """
        me=np.eye(12,dtype=float)*(self.rho*self.A*self.L/2.0)
        # 质量矩阵的对角线元素
        me[(3,9),(3,9)]*=self.J/self.A
        me[(4,5,10,11),(4,5,10,11)]=0.0
        return self.t.T@me@self.t
    @property
    def D(self):
        """
        返回单元的弹性模量E
        
        returns:
            float: 弹性模量E
        """
        return self.E
    @property
    def dofNum(self):
        """
        获取单元的自由度数量
        
        returns:
            int: 自由度数量
        """
        return 12
    @property
    def NodesList(self):
        """
        获取单元的节点列表
        
        returns:
            List[Node]: 节点列表
        """
        return [self.node1, self.node2]
    @property
    def ElemDofIndexs(self):
        """
        获取单元的全局自由度索引列表(0-based)
        
        returns:
            List[int]: 自由度索引列表
        """
        return dofInds(self.node1.label, self.dofNum) + dofInds(self.node2.label, self.dofNum)
        
class F2D2_EB(elementBase):
    """
    平面框架单元,EB后缀代表采用欧拉伯努利梁假设,节点自由度:UX,UY,URZ;
    
    可承受载荷: 轴力F1,F2,弯矩Mz,不支持单元载荷.
    
    单元定义需要:A,E,rho,节点坐标,Izz;(Izz为截面惯性矩);
    
    单元公式见项目文档;
    """
    VTK_TYPE = 3
    def __init__(self, label:int,
                    node1:Node,
                    node2:Node,
                    E:float=0,A:float=0,Izz:float=0,rho:float=0.0):
        super().__init__(label, 3)
        self.node1=node1
        self.node2=node2
        self.E=E
        self.A=A
        self.Izz=Izz
        self.rho=rho
    @property
    def get_connectivity(self):
        """
        获取单元节点编号
        
        returns:
            List[int]: 节点编号列表
        """
        return [self.node1.label,self.node2.label]
    @property
    def L(self)->float:
        """
        计算单元长度
        
        returns:
            float: 单元长度
        """
        return np.sqrt((self.node2.x-self.node1.x)**2+(self.node2.y-self.node1.y)**2)
    @property
    def t(self):
        """
        计算坐标变换矩阵
        
        returns:
            np.ndarray: 坐标变换矩阵(6,6)
        """
        # 计算局部x轴,并归一化
        axis_x=self.node2.coordinate-self.node1.coordinate
        axis_x[2]=0.0  # 保证z坐标为0
        axis_x/=np.linalg.norm(axis_x)
        # x轴绕全局Z轴按右手定则旋转90度得到局部y轴
        axis_y=np.array([-axis_x[1], axis_x[0], 0.0])
        return np.array([[axis_x[0], axis_x[1], axis_x[2], 0, 0, 0],
                            [axis_y[0], axis_y[1], axis_y[2], 0, 0, 0],
                            [0,0,1,0,0,0],
                            [0,0,0,axis_x[0], axis_x[1], axis_x[2]],
                            [0,0,0,axis_y[0], axis_y[1], axis_y[2]],
                            [0,0,0,0,0,1]],dtype=float)
    @property
    def Ke(self)->np.ndarray:
        """
        全局坐标系下的单元刚度矩阵
        
        returns:
            np.ndarray: 单元刚度矩阵(6,6)
        """
        ke=np.zeros((6,6),dtype=float)
        # 装配 Axial Displacements 的刚度矩阵部分
        ke[0,0]+=(self.E*self.A)/self.L
        ke[3,3]+=ke[0,0]
        ke[0,3]+=-1*(self.E*self.A)/self.L
        ke[3,0]+=ke[0,3]
        # 加入 Bending Displacements in the Plane xy 的刚度矩阵部分
        kxy=np.array(([[12, 6*self.L,-12, 6*self.L],
                        [6*self.L, 4*self.L**2,-6*self.L, 2*self.L**2],
                        [-12,-6*self.L,12,-6*self.L],
                        [6*self.L, 2*self.L**2,-6*self.L, 4*self.L**2]]))*(self.E*self.Izz/(self.L**3))
        ke[np.ix_([1,2,4,5],[1,2,4,5])]+=kxy
        return self.t.T@ke@self.t
    @property
    def Me(self)->np.ndarray:
        """
        全局坐标系下的集中单元质量矩阵
        
        returns:
            np.ndarray: 单元质量矩阵(6,6)
        """
        me=np.eye(6,dtype=float)*(self.rho*self.A*self.L/2.0)
        me[(2,5),(2,5)]=0.0  # 取消URZ的质量
        return me
    @property
    def D(self):
        """
        返回单元的弹性模量E
        returns:
            float: 弹性模量E
        """
        return self.E
    @property
    def dofNum(self)->int:
        return 6
    @property
    def NodesList(self)->List[Node]:
        return [self.node1,self.node2]
    @property
    def ElemDofIndexs(self)->List[int]:
        return dofInds(self.node1.label,self.NodeDofNum)+dofInds(self.node2.label,self.NodeDofNum)
    # def ProcessELoad2EqNodeForce(self, load:loadBase):
    #     """
    #     根据载荷类型将单元载荷转换为等效节点力
    #     """
    #     if isinstance(load,DloadOnBeam):
    #         return self.DloadonBeam2EqNodeForce(load)
    #     elif isinstance(load,DloadOnTruss):
    #         return self.DloadOnTruss2EqNodeForce(load)
    # def DloadonBeam2EqNodeForce(self, load:DloadOnBeam)->np.ndarray:
    #     """将横向均布荷载荷DloadOnBeam转换为等效节点荷载"""
    #     pe=(load.LoadValue/12.0)*np.array([[0],[6*self.length],[self.length*self.length],[0],[6*self.length],[-1.0*self.length*self.length]])
    #     EqForce=self.T.T.dot(pe)
    #     if EqForce.shape!=(6,1):
    #         raise ValueError("EqForce shape is not (6,1)!")
    #     return EqForce
    # def DloadOnTruss2EqNodeForce(self, load:DloadOnTruss)->np.ndarray:
    #     """将均布荷载荷DloadOnTruss转换为等效节点荷载"""
    #     pe=np.array([[1],[0],[0],[1],[0],[0]])*(0.5*load.LoadValue*self.length)
    #     eqforce=(self.T.T).dot(pe)
    #     if eqforce.shape!=(6,1):
    #         raise ValueError("EqForce shape is not (6,1)!")
    #     return eqforce

